#include "MyNpzScoreWriter.hh"

#include <G4Types.hh>
#include <vector>

#include "G4StatDouble.hh"
#include "G4VScoreWriter.hh"  // 先让父类写 CSV
#include "G4VScoringMesh.hh"
#include "cnpy.h"
// #include "G4SystemOfUnits.hh"

#define _MY_WRITE_RLTV_ERR 1

void MyNpzScoreWriter::DumpQuantityToFile(const G4String& psName, const G4String& fileName, const G4String& option)
{
    ///==================================== G4 11.3.2
    using ScoreMap = std::map<G4int, G4StatDouble*>;

    ///======================================== 保留官方 CSV 输出, for debug
    G4VScoreWriter::DumpQuantityToFile(psName, fileName, option);

    ///======================================== 构造 .npy 文件名
    G4String npyName = fileName;
    std::size_t dot = npyName.rfind('.');
    if (dot != std::string::npos) npyName.erase(dot);
    npyName += ".npy";

    if (!fScoringMesh) return;
    auto scoreMap = fScoringMesh->GetScoreMap();
    auto it = scoreMap.find(psName);
    if (it == scoreMap.end()) return;

    ScoreMap* const score = it->second->GetMap();  // 返回的是 G4StatDouble*

    /// 3D 缓冲区, 存"均值"
    std::vector<G4double> buf(fNMeshSegments[0] * fNMeshSegments[1] * fNMeshSegments[2], 0.0);

#if _MY_WRITE_RLTV_ERR
    /// 可选: 同时写一份"相对误差"
    std::vector<G4double> rltvErr(buf.size(), 0.0);
#endif

    const G4double unitValue = fScoringMesh->GetPSUnitValue(psName);

    for (int ix = 0; ix < fNMeshSegments[0]; ++ix)
    {
        for (int iy = 0; iy < fNMeshSegments[1]; ++iy)
        {
            for (int iz = 0; iz < fNMeshSegments[2]; ++iz)
            {
                G4int idx = this->GetIndex(ix, iy, iz);
                auto vit = score->find(idx);
                if (vit != score->end() && vit->second)
                {
                    G4StatDouble* const stat = vit->second;
                    const G4double meanVal = stat->mean() / unitValue;
                    buf[idx] = meanVal;

#if _MY_WRITE_RLTV_ERR
                    // 相对误差(rms/mean), 避免除 0; root mean square
                    const G4double rmsVal = stat->rms();
                    if (meanVal != 0.0) rltvErr[idx] = rmsVal / meanVal;
#endif
                }
            }
        }
    }

    // 写均值
    cnpy::npy_save(npyName, buf.data(), {size_t(fNMeshSegments[0]), size_t(fNMeshSegments[1]), size_t(fNMeshSegments[2])}, "w");

#if _MY_WRITE_RLTV_ERR
    // 如需误差, 再写一份 *_err.npy
    cnpy::npy_save(npyName + "_err", rltvErr.data(), {size_t(fNMeshSegments[0]), size_t(fNMeshSegments[1]), size_t(fNMeshSegments[2])}, "w");
#endif
}